import { LoadingOutlined, TagsOutlined, UserOutlined } from "@ant-design/icons";
import { ThoughtChain, XRequest, Bubble } from "@ant-design/x";
import { ThoughtChainItem } from "@ant-design/x/es/components/thought-chain";
import {
  Button,
  Descriptions,
  Flex,
  Input,
  Splitter,
  Typography,
} from "antd";
import React, { useRef, useState } from "react";
import { produce } from "immer";


const { Paragraph } = Typography;

/**
 * 🔔 Please replace the BASE_URL, PATH, MODEL, API_KEY with your own values.
 */

const BASE_URL = `${import.meta.env.VITE_BASE_URL}/user/chat/new/generate`;

/**
 * 🔔 The MODEL is fixed in the current request, please replace it with your BASE_UR and MODEL
 */

const MODEL = "Qwen3-8B";

/**
 * 🔔 the API_KEY is a placeholder indicator interface that has a built-in real API_KEY
 */

const API_KEY = "Bearer sk-xxxxxxxxxxxxxxxxxxxx";
const auth_header = "Basic cm9vdDpyb290";
const authorization =
  "bearer eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJpc3MiOiJodHRwOlwvXC8xMDEuMjAwLjE1Ny4yMjA6NjY3XC9hcGlcL2F1dGhcL2xvZ2luIiwiaWF0IjoxNzQ5MDkyMzY3LCJleHAiOjE3NDk5NTYzNjcsIm5iZiI6MTc0OTA5MjM2NywianRpIjoiODFLSEdyUXgyMFRiaVB1NCIsInN1YiI6NSwicHJ2IjoiZjY0ZDQ4YTZjZWM3YmRmYTdmYmY4OTk0NTRiNDg4YjNlNDYyNTIwYSJ9.xAbDdCSONLWAGptooXDKiQBOTyz7SiFSPdzUH2xP538";
// 自定义fetch函数，添加认证头部
const customFetch = async (input, init) => {
  const customHeaders = {
    "Content-Type": "application/json",
    auth_header: auth_header, // 替换为你的实际 auth_header 值
    authorization: authorization, // 使用上面定义的API_KEY或替换为你的实际值
    ...init?.headers, // 保留原有的headers
  };

  const customInit = {
    ...init,
    headers: customHeaders,
  };

  return fetch(input, customInit);
};

const exampleRequest = XRequest({
  baseURL: BASE_URL,
  // model: MODEL,
  fetch: customFetch, // 使用自定义的fetch函数
  /** 🔥🔥 Its dangerously! */
});
const roles = {
  assistant: {
    placement: "start",
    avatar: { icon: <UserOutlined />, style: { background: "#fde3cf" } },
  },
  user: {
    placement: "end",
    avatar: { icon: <UserOutlined />, style: { background: "#87d068" } },
  },
};

const App = () => {
  const [status, setStatus] = useState();
  const [chatMessageList, setChatMessageList] = useState([]);
  const [thoughtChainStatus, setThoughtChainStatus] =
    useState();
  const [lines, setLines] = useState([]);
  const [questionText, setQuestionText] = useState("hello, who are u?");
  const abortController = useRef(null);

  const request = async () => {
    setStatus("pending");
    setLines([]);
    
    // 添加用户消息到聊天列表
    const userMessageId = Date.now();
    const assistantMessageId = userMessageId + 1;
    
    setChatMessageList(
      produce((draft) => {
        draft.push(
          {
            id: userMessageId,
            message: {
              type: 0, // 用户消息
              message: questionText,
            },
          },
          {
            id: assistantMessageId,
            message: {
              type: 1, // 助手消息
              message: "",
            },
          }
        );
      })
    );

    await exampleRequest.create(
      {
        messages: [{ role: "user", content: questionText }],
        stream: true,
        chat_id: 6,
      },
      {
        onSuccess: () => {
          console.log("onSuccess");
          setStatus("success");
          setThoughtChainStatus("success");
        },
        onError: (error) => {
          if (error.name === "AbortError") {
            setStatus("abort");
          }
          setThoughtChainStatus("error");
        },
        onUpdate: (msg) => {
          console.log("msg", msg);
          
          // 检查是否完成
          if (msg.data.includes("[DONE]")) {
            setStatus("success");
            setThoughtChainStatus("success");
            return;
          }
          
          try {
            const _msg = JSON.parse(msg.data);
            const resultMessage = _msg?.choices?.[0]?.delta?.content;
            console.log("resultMessage", resultMessage);
            
            if (resultMessage) {
              // 使用 immer 更新最后一个助手消息的内容
              setChatMessageList(
                produce((draft) => {
                  const lastMessage = draft[draft.length - 1];
                  if (lastMessage?.message.type === 1) {
                    lastMessage.message.message += resultMessage;
                  }
                })
              );
            }
          } catch (error) {
            console.error("解析消息失败:", error);
          }

          setLines((pre) => [...pre, msg]);
        },
        onStream: (controller) => {
          abortController.current = controller;
        },
      }
    );
  };

  const abort = () => {
    abortController?.current?.abort?.();
    setStatus("abort");
    setThoughtChainStatus("error");
  };

  return (
    <div>
      <Bubble.List
        roles={roles}
        style={{ maxHeight: 300 }}
        items={chatMessageList.map(({ id, message }) => ({
          key: id,
          role: message.type === 1 ? "assistant" : "user",
          content: message.message,
        }))}
      />
      <Splitter>
        <Splitter.Panel style={{ height: 300 }}>
          <Splitter layout="vertical">
            <Splitter.Panel style={{ margin: "0 16px" }}>
              <Flex gap="large" vertical>
                <Input
                  value={questionText}
                  onChange={(e) => {
                    setQuestionText(e.target.value);
                  }}
                />
                <Flex gap="small">
                  <Button
                    type="primary"
                    disabled={status === "pending"}
                    onClick={request}
                  >
                    Request
                  </Button>
                  <Button
                    type="primary"
                    disabled={status !== "pending"}
                    onClick={abort}
                  >
                    Request Abort
                  </Button>
                </Flex>
              </Flex>
            </Splitter.Panel>
            <Splitter.Panel style={{ margin: 16 }}>
              <Paragraph>{lines.length > 0 && JSON.stringify(lines)}</Paragraph>
            </Splitter.Panel>
          </Splitter>
        </Splitter.Panel>
        <Splitter.Panel style={{ marginLeft: 16 }}>
          <ThoughtChain
            items={[
              {
                title: "Request Log",
                status: thoughtChainStatus,
                icon:
                  status === "pending" ? <LoadingOutlined /> : <TagsOutlined />,
                description: `request ${status}`,
                content: (
                  <Descriptions column={1}>
                    <Descriptions.Item label="Status">
                      {status || "-"}
                    </Descriptions.Item>
                    <Descriptions.Item label="Update Times">
                      {lines.length}
                    </Descriptions.Item>
                  </Descriptions>
                ),
              },
            ]}
          />
        </Splitter.Panel>
      </Splitter>
    </div>
  );
};

export default App;
